from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from typing import Dict, Set
import json
import uuid
import os
import asyncio

app = FastAPI()

# 更新CORS配置，使其更灵活
origins = [
    "http://121.41.104.175:8001",
    "http://121.41.104.175:8000",
    "http://localhost:8001",
    "http://localhost:8000",
    "http://127.0.0.1:8001",
    "http://127.0.0.1:8000",
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["*"],
    max_age=3600,
)

# Mount static files with proper MIME types
app.mount("/static", StaticFiles(directory="static"), name="static")
app.mount("/assets", StaticFiles(directory="static/assets"), name="assets")


# Serve index.html for all non-API routes
@app.get("/")
async def serve_root():
    return FileResponse("static/index.html")



# Store active connections and room data
class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, Dict[str, WebSocket]] = {}
        self.room_data: Dict[str, list] = {}

    async def connect(self, websocket: WebSocket, room_id: str, client_id: str):
        try:
            await websocket.accept()
            if room_id not in self.active_connections:
                self.active_connections[room_id] = {}
                self.room_data[room_id] = []
            self.active_connections[room_id][client_id] = websocket
            print(f"Active connections in room {room_id}: {len(self.active_connections[room_id])}")
        except Exception as e:
            print(f"Error in connect: {str(e)}")
            raise

    def disconnect(self, room_id: str, client_id: str):
        if room_id in self.active_connections:
            self.active_connections[room_id].pop(client_id, None)
            if not self.active_connections[room_id]:
                self.active_connections.pop(room_id)
                self.room_data.pop(room_id, None)

    async def send_initial_state(self, websocket: WebSocket, room_id: str):
        """发送房间的初始状态给新连接的客户端"""
        if room_id not in self.room_data:
            return

        room_state = self.room_data[room_id]
        if not room_state:
            return

        try:
            # 发送初始化开始信号
            await websocket.send_json({
                "type": "init_start",
                "message": "开始发送房间状态"
            })

            # 分块发送状态数据
            chunk_size = 50  # 减小块大小以提高性能
            for i in range(0, len(room_state), chunk_size):
                chunk = room_state[i:i + chunk_size]
                await websocket.send_json({
                    "type": "init_chunk",
                    "data": chunk,
                    "chunk_index": i // chunk_size
                })
                await asyncio.sleep(0.1)  # 添加小延迟避免消息堆积

            # 发送初始化完成信号
            await websocket.send_json({
                "type": "init_complete",
                "total_chunks": len(room_state) // chunk_size + 1
            })
        except Exception as e:
            print(f"Error sending initial state: {str(e)}")
            raise

    async def broadcast_to_room(self, message: str, room_id: str, sender_id: str):
        if room_id in self.active_connections:
            disconnected_clients = []
            for client_id, connection in self.active_connections[room_id].items():
                if client_id != sender_id:
                    try:
                        await connection.send_text(message)
                    except Exception as e:
                        print(f"Error sending to client {client_id}: {str(e)}")
                        disconnected_clients.append(client_id)
            
            # 清理断开的连接
            for client_id in disconnected_clients:
                self.disconnect(room_id, client_id)

    def save_room_state(self, room_id: str, state_data: dict):
        """保存房间状态"""
        if room_id in self.room_data:
            self.room_data[room_id].append(state_data)
            # 可以选择限制历史记录的大小
            if len(self.room_data[room_id]) > 1000:  # 限制最大历史记录数
                self.room_data[room_id] = self.room_data[room_id][-1000:]


manager = ConnectionManager()


@app.get("/create-room")
async def create_room():
    print("收到创建房间请求")  # 添加日志
    try:
        room_id = str(uuid.uuid4())
        print(f"创建了新房间: {room_id}")  # 添加日志
        return {"room_id": room_id}
    except Exception as e:
        print(f"创建房间时出错: {str(e)}")  # 添加错误日志
        raise


@app.websocket("/ws/{room_id}")
async def websocket_endpoint(websocket: WebSocket, room_id: str):
    client_id = str(uuid.uuid4())
    try:
        print(f"新客户端 {client_id} 尝试连接到房间 {room_id}")
        await manager.connect(websocket, room_id, client_id)
        print(f"客户端 {client_id} 已连接到房间 {room_id}")

        # 发送连接成功的确认消息
        await websocket.send_json({
            "type": "connection_established",
            "client_id": client_id,
            "room_id": room_id
        })

        while True:
            data = await websocket.receive_text()
            message_data = json.loads(data)
            message_type = message_data.get('type')
            print(f"收到来自客户端 {client_id} 的消息: {message_type}")

            # 处理客户端准备就绪的消息
            if message_type == "client_ready":
                try:
                    # 发送房间的初始状态
                    await manager.send_initial_state(websocket, room_id)
                    await websocket.send_json({
                        "type": "server_ack",
                        "message": "服务器已确认客户端就绪状态"
                    })
                except Exception as e:
                    print(f"发送初始状态时出错: {str(e)}")
                continue

            # 保存绘画状态
            if message_type == "draw":
                try:
                    manager.save_room_state(room_id, message_data)
                    print(f"已保存房间 {room_id} 的绘画状态")
                    # 发送保存确认
                    await websocket.send_json({
                        "type": "state_saved",
                        "timestamp": message_data.get('timestamp')
                    })
                except Exception as e:
                    print(f"保存状态时出错: {str(e)}")
                    await websocket.send_json({
                        "type": "error",
                        "message": "保存状态失败"
                    })

            # 保存回退状态
            if message_type == "undo":
                try:
                    manager.save_room_state(room_id, message_data)
                    print(f"已保存房间 {room_id} 的回退状态")
                    # 发送保存确认
                    await websocket.send_json({
                        "type": "state_saved",
                        "timestamp": message_data.get('timestamp')
                    })
                except Exception as e:
                    print(f"保存回退状态时出错: {str(e)}")
                    await websocket.send_json({
                        "type": "error",
                        "message": "保存回退状态失败"
                    })

            # 处理清除画布操作
            if message_type == "clear":
                try:
                    # 清除房间的历史状态
                    if room_id in manager.room_data:
                        manager.room_data[room_id] = []
                    
                    # 保存清除状态
                    manager.save_room_state(room_id, message_data)
                    print(f"已清除房间 {room_id} 的画布状态")
                    
                    # 发送保存确认
                    await websocket.send_json({
                        "type": "state_saved",
                        "message": "画布已清除",
                        "timestamp": message_data.get('timestamp')
                    })
                except Exception as e:
                    print(f"清除画布状态时出错: {str(e)}")
                    await websocket.send_json({
                        "type": "error",
                        "message": "清除画布状态失败"
                    })

            # 广播消息给房间内的其他用户
            await manager.broadcast_to_room(data, room_id, client_id)
            print(f"已广播消息到房间 {room_id}")

    except WebSocketDisconnect:
        print(f"客户端 {client_id} 与房间 {room_id} 断开连接")
        manager.disconnect(room_id, client_id)
        try:
            await manager.broadcast_to_room(
                json.dumps({
                    "type": "disconnect",
                    "client_id": client_id,
                    "timestamp": str(uuid.uuid4())
                }),
                room_id,
                client_id
            )
        except Exception as e:
            print(f"广播断开连接消息时出错: {str(e)}")
    except Exception as e:
        print(f"WebSocket连接出错: {str(e)}")
        manager.disconnect(room_id, client_id)
        try:
            await websocket.send_json({
                "type": "error",
                "message": str(e)
            })
        except:
            pass


@app.get("/health")
async def health_check():
    return {"status": "healthy"}


if __name__ == "__main__":
    import uvicorn
    
    uvicorn.run(
        app,
        host="0.0.0.0",  # 允许所有IP访问
        port=8001,
        log_level="info"  # 添加日志级别以便调试
    )
